from flask import Flask, request, jsonify
import threading
# import os
import config
from pathlib import Path
# import numpy as np
from obs import ObsClient
from groundingdino.util.inference import load_model
from pcdet.config import cfg, cfg_from_yaml_file
from pcdet.models import build_network
# from pcdet.utils import common_utils
from ultralytics import YOLO
from background_task import DemoDataset, background_task_data_annotation, background_task_data_mining, data_path_dict
import logging
import logging.handlers
import sys
# logger = common_utils.create_logger()
logger = logging.getLogger("my_log")
logger.setLevel(logging.INFO)

console_handler = logging.StreamHandler(sys.stdout)
console_handler.setLevel(logging.INFO)

file_handler = logging.handlers.RotatingFileHandler("my_log.log", maxBytes=10**6, backupCount=3)
file_handler.setLevel(logging.INFO)

formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
console_handler.setFormatter(formatter)
file_handler.setFormatter(formatter)

logger.addHandler(console_handler)
logger.addHandler(file_handler)

try:
    #load model only once
    groundingdino_model = load_model("./img_detection_groundingdino/groundingdino/config/GroundingDINO_SwinT_OGC.py", "./img_detection_groundingdino/weights/groundingdino_swint_ogc.pth")
    groundingdino_model.cuda()
    groundingdino_model.eval()
    logger.info("groundingdino model load done")
    cfg_from_yaml_file("./pcd_detection/tools/cfgs/nuscenes_models/cbgs_voxel0075_voxelnext.yaml", cfg)
    demo_dataset = DemoDataset(dataset_cfg=cfg.DATA_CONFIG, class_names=cfg.CLASS_NAMES, training=False, root_path=Path("./pcd_detection/tools/000690_pandarat128_2372.pcd"), ext=".pcd", logger=logger)
    voxelnext_model = build_network(model_cfg=cfg.MODEL, num_class=len(cfg.CLASS_NAMES), dataset=demo_dataset)
    voxelnext_model.load_params_from_file(filename="./pcd_detection/voxelnext_nuscenes_kernel1.pth", logger=logger, to_cpu=True)
    voxelnext_model.cuda()
    voxelnext_model.eval()
    logger.info("voxelnext model to gpu done")
    yolo_model = YOLO("./img_detection_yolo/weight/v8x_best_0717.pt")
    yolo_model.cuda()
    logger.info("yolo model to gpu done")
    obs_client_annotation = ObsClient(access_key_id=config.annotation_ak, secret_access_key=config.annotation_sk, server=config.annotation_endpoint)
    obs_client_datacenter = ObsClient(access_key_id=config.ak, secret_access_key=config.sk, server=config.endpoint)
    logger.info("obs_client init done")
except Exception as e:
    logger.error("init model or obs error:{}".format(e))

app = Flask(__name__)

@app.route('/dataAnnotation', methods=['POST'])
def data_annotation():
    # check parameters
    try:
        #in json, there is datasetId, datasetType, datasetPath, imageConfTh, pcdConfTh
        data = request.get_json()
        logger.info("/dataAnnotation receive data:{}".format(data))
        # check json file content
        datasetId = data.get("datasetId")
        if not isinstance(datasetId, str):
            error_message = "datasetId {} type check failed, need str".format(datasetId)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        
        datasetType = data.get("datasetType")
        if not isinstance(datasetType, str):
            error_message = "datasetType {} type check failed, need str".format(datasetType)
            logger.error(error_message)
            return jsonify({"error": error_message}), 400
        if datasetType not in ["1", "2", "3"]:
            error_message = "datasetType {} value check failed, need '1' for image annotation, '2' for pcd annotation, '3' for joint annotation".format(datasetType)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        
        datasetPath = data.get("datasetPath")
        if not isinstance(datasetPath, str):
            error_message = "datasetPath {} type check failed, need str".format(datasetPath)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        if not datasetPath.endswith("/"):
            datasetPath = datasetPath + "/" 
        
        callbackUrl = data.get("callbackUrl")
        if not isinstance(callbackUrl, str):
            error_message = "callbackUrl {} type check failed, need str".format(callbackUrl)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        
        imageConfTh = data.get("imageConfTh")
        if not isinstance(imageConfTh, float):
            error_message = "imageConfTh {} type check failed, need float".format(imageConfTh)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        if imageConfTh < 0.000001 or imageConfTh > 0.999999: 
            error_message = "imageConfTh {} value check failed, need [0.01, 0.99]".format(imageConfTh)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        
        pcdConfTh = data.get("pcdConfTh")
        if not isinstance(pcdConfTh, float):
            error_message = "pcdConfTh {} type check failed, need float".format(pcdConfTh)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        if pcdConfTh < 0.00001 or pcdConfTh > 0.999999: 
            error_message = "pcdConfTh {} value check failed, need [0.01, 0.99]".format(pcdConfTh)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        
        # check bucket download
        response = obs_client_annotation.listObjects(config.annotation_bucket_name, prefix=datasetPath)
        if response.status < 300:
            logger.info('{} listObjects success'.format(config.annotation_bucket_name + "/" + datasetPath))
        else:
            error_message = '{} listObjects failed'.format(config.annotation_bucket_name + "/" + datasetPath)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        objects = []
        for obj in response.body.contents:
            if obj.key.endswith(".zip"):
                objects.append(obj.key)
        if len(objects) == 0:
            error_message = 'no zip file in {}'.format(config.bucket_name + "/" + datasetPath)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        else:
            logger.info('{} zip files in {}'.format(len(objects), config.bucket_name + "/" + datasetPath))

       # check bucket upload 
        object_key = data_path_dict["pre_annotation_upload_path"] + datasetPath
        response = obs_client_annotation.putFile(config.annotation_bucket_name, object_key+"test_upload.json", file_path="./test_upload.json")
        if response.status < 300:
            logger.info('{} putFile test_upload success'.format(config.annotation_bucket_name + "/" + object_key))
        else:
            error_message = '{} putFile test_upload failed'.format(config.annotation_bucket_name + "/" + object_key)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        #start annotation
        threading.Thread(target=background_task_data_annotation, args=(voxelnext_model, yolo_model, obs_client_annotation, data, logger)).start()
        logger.info("background_task_data_annotation thread start")
        # notify platform that post checked ok
        return jsonify({"code": "200", "mesage": "data checked ok"}), 200
    except Exception as e:
        # 400 for known error, 600 for unknown error
        logger.error("/dataAnnotation meet error {}".format(e))
        return {"code": "600", "mesage": e}, 600


@app.route('/dataMining', methods=['POST'])
def data_mining():
    # check parameters
    try:
        #in json, there is taskId, datasetPath, startFrameNo, endFrameNo, labelsId, prompts, intervalLabelsId, intervalPrompts, boxTh, textTh, callbackUrl
        data = request.get_json()
        logger.info("/dataMining receive data:{}".format(data))
        # check json file content
        taskId = data.get("taskId")
        if not isinstance(taskId, str):
            error_message = "taskId {} type check failed, need str".format(taskId)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        
        datasetPath = data.get("datasetPath")
        if not isinstance(datasetPath, str):
            error_message = "datasetPath {} type check failed, need str".format(datasetPath)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        if not datasetPath.endswith("/"):
            datasetPath = datasetPath + "/" 
        
        callbackUrl = data.get("callbackUrl")
        if not isinstance(callbackUrl, str):
            error_message = "callbackUrl {} type check failed, need str".format(callbackUrl)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        
        labelsId = data.get("labelsId")
        if not isinstance(labelsId, str):
            error_message = "labelsId {} type check failed, need str".format(labelsId)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        
        prompts = data.get("prompts")
        if not isinstance(prompts, str):
            error_message = "prompts {} type check failed, need str".format(prompts)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        if "." not in prompts:
            error_message = "prompts {} value check failed, must exist at least one '.'".format(prompts)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400

        labelsId_len = len(labelsId.split(","))
        prompts_len = len(prompts.split(".")) - 1 #"cat . dog . cow ."
        if not labelsId_len == prompts_len:
            error_message = "prompts {} len:{} and labelsId {} len:{} check failed, need to be same".format(prompts, prompts_len, labelsId, labelsId_len)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        
        startFrameNo = data.get("startFrameNo")
        if not isinstance(startFrameNo, int):
            error_message = "startFrameNo {} type check failed, need int".format(startFrameNo)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        if startFrameNo <= 0:
            error_message = "startFrameNo {} value check failed, need > 0".format(startFrameNo)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        
        endFrameNo = data.get("endFrameNo")
        if not isinstance(endFrameNo, int):
            error_message = "endFrameNo {} type check failed, need int".format(endFrameNo)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        if endFrameNo <= 0:
            error_message = "endFrameNo {} value check failed, need > 0".format(endFrameNo)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        
        if not endFrameNo > startFrameNo:
            error_message = "endFrameNo {} and startFrameNo {} value check failed, endFrameNo should be bigger".format(endFrameNo, startFrameNo)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        
        boxTh = data.get("boxTh")
        if not isinstance(boxTh, float):
            error_message = "boxTh {} type check failed, need float".format(boxTh)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        if boxTh < 0.000001 or boxTh > 0.999999: 
            error_message = "boxTh {} value check failed, need [0.01, 0.99]".format(boxTh)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        
        text_conf_th = data.get("textTh")
        if not isinstance(text_conf_th, float):
            error_message = "text_conf_th {} type check failed, need float".format(text_conf_th)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        if text_conf_th < 0.00001 or text_conf_th > 0.999999: 
            error_message = "text_conf_th {} value check failed, need [0.01, 0.99]".format(text_conf_th)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        
        # check bucket download
        response = obs_client_datacenter.listObjects(config.bucket_name, prefix=datasetPath)#
        if response.status < 300:
            logger.info('{} listObjects success'.format(config.bucket_name + "/" + datasetPath))
        else:
            error_message = '{} listObjects failed'.format(config.bucket_name + "/" + datasetPath)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        objects = []
        for obj in response.body.contents:
            if obj.key.endswith(".zip"):
                objects.append(obj.key)
        if len(objects) == 0:
            error_message = 'no zip file in {}'.format(config.bucket_name + "/" + datasetPath)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        else:
            logger.info('{} zip files in {}'.format(len(objects), config.bucket_name + "/" + datasetPath))

        # check bucket upload 
        object_key = data_path_dict["mining_upload_path"] + datasetPath
        response = obs_client_datacenter.putFile(config.bucket_name, object_key+"test_upload.json", file_path="./test_upload.json")
        if response.status < 300:
            logger.info('{} putFile test_upload success'.format(config.bucket_name + "/" + object_key))
        else:
            error_message = '{} putFile test_upload failed'.format(config.bucket_name + "/" + object_key)
            logger.error(error_message)
            return jsonify({"code": "400", "mesage": error_message}), 400
        #data type and value checked ok, download and upload file checked ok, start mining thread
        threading.Thread(target=background_task_data_mining, args=(groundingdino_model, obs_client_datacenter, data, logger)).start()
        logger.info("background_task_data_mining thread start")
        # notify platform that post checked ok
        return jsonify({"code": "200", "mesage": "data checked ok"}), 200
    except Exception as e:
        # 400 for known error, 600 for unknown error
        logger.error("/dataMining meet error {}".format(e))
        return {"code": "600", "mesage": e}, 600


if __name__ == '__main__':
    try:
        app.run(debug=True, use_reloader=False, host=config.host, port=config.port)#threaded=False, processes=5, debug=False
    except Exception as e:
        logger.error("app run error:{}".format(e))
    
